#!/bin/bash

# AIQuant 健康检查脚本
# 用于检查AIQuant量化投研系统的健康状态

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 健康状态
HEALTH_STATUS=0  # 0=健康, 1=警告, 2=严重

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[OK]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
    HEALTH_STATUS=1
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
    HEALTH_STATUS=2
}

# 检查API服务健康状态
check_api_health() {
    local endpoint="http://localhost:8000/health"
    local timeout=10
    
    if curl -f -s --max-time $timeout "$endpoint" > /dev/null 2>&1; then
        local response=$(curl -s --max-time $timeout "$endpoint")
        local status=$(echo "$response" | python3 -c "import sys, json; print(json.load(sys.stdin).get('status', 'unknown'))" 2>/dev/null || echo "unknown")
        
        if [ "$status" = "healthy" ]; then
            log_success "API服务健康状态正常"
        else
            log_warning "API服务状态异常: $status"
        fi
    else
        log_error "API服务健康检查失败"
    fi
}

# 检查Web服务健康状态
check_web_health() {
    local endpoint="http://localhost:8501"
    local timeout=10
    
    if curl -f -s --max-time $timeout "$endpoint" > /dev/null 2>&1; then
        log_success "Web服务响应正常"
    else
        log_error "Web服务无响应"
    fi
}

# 检查数据库连接
check_database_health() {
    python3 -c "
import sys
sys.path.append('.')
try:
    from aiquant.database.connection import DatabaseManager
    db = DatabaseManager()
    
    # 测试连接
    db.test_connection()
    
    # 检查表是否存在
    with db.get_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(\"SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = 'public'\")
        table_count = cursor.fetchone()[0]
        
        if table_count > 0:
            print(f'数据库连接正常，共有 {table_count} 个表')
        else:
            print('数据库连接正常但无数据表')
            sys.exit(1)
            
except ImportError:
    print('无法导入数据库模块')
    sys.exit(2)
except Exception as e:
    print(f'数据库连接失败: {e}')
    sys.exit(2)
" 2>/dev/null

    case $? in
        0)
            log_success "数据库健康状态正常"
            ;;
        1)
            log_warning "数据库连接正常但结构异常"
            ;;
        2)
            log_error "数据库连接失败"
            ;;
    esac
}

# 检查Redis连接
check_redis_health() {
    if command -v redis-cli &> /dev/null; then
        if redis-cli ping > /dev/null 2>&1; then
            # 检查内存使用情况
            local memory_info=$(redis-cli info memory | grep "used_memory_human" | cut -d: -f2 | tr -d '\r')
            log_success "Redis连接正常，内存使用: $memory_info"
            
            # 检查连接数
            local connected_clients=$(redis-cli info clients | grep "connected_clients" | cut -d: -f2 | tr -d '\r')
            if [ "$connected_clients" -gt 100 ]; then
                log_warning "Redis连接数较高: $connected_clients"
            fi
        else
            log_error "Redis连接失败"
        fi
    else
        log_warning "Redis CLI未安装，无法检查Redis状态"
    fi
}

# 检查系统资源
check_system_resources() {
    # 检查CPU使用率
    local cpu_usage=$(top -bn1 | grep "Cpu(s)" | awk '{print $2}' | cut -d'%' -f1 | cut -d',' -f1)
    if (( $(echo "$cpu_usage > 80" | bc -l) )); then
        log_warning "CPU使用率过高: ${cpu_usage}%"
    else
        log_success "CPU使用率正常: ${cpu_usage}%"
    fi
    
    # 检查内存使用率
    local mem_info=$(free | grep Mem)
    local total_mem=$(echo $mem_info | awk '{print $2}')
    local used_mem=$(echo $mem_info | awk '{print $3}')
    local mem_usage=$(echo "scale=1; $used_mem * 100 / $total_mem" | bc)
    
    if (( $(echo "$mem_usage > 85" | bc -l) )); then
        log_warning "内存使用率过高: ${mem_usage}%"
    else
        log_success "内存使用率正常: ${mem_usage}%"
    fi
    
    # 检查磁盘使用率
    local disk_usage=$(df / | tail -1 | awk '{print $5}' | cut -d'%' -f1)
    if [ "$disk_usage" -gt 90 ]; then
        log_error "磁盘使用率过高: ${disk_usage}%"
    elif [ "$disk_usage" -gt 80 ]; then
        log_warning "磁盘使用率较高: ${disk_usage}%"
    else
        log_success "磁盘使用率正常: ${disk_usage}%"
    fi
}

# 检查进程状态
check_processes() {
    local processes=("api" "web" "data_collector" "market_monitor" "anomaly_detector")
    
    for process in "${processes[@]}"; do
        local pid_file=".${process}.pid"
        
        if [ -f "$pid_file" ]; then
            local pid=$(cat "$pid_file")
            if ps -p $pid > /dev/null 2>&1; then
                # 检查进程资源使用情况
                local cpu_usage=$(ps -p $pid -o %cpu --no-headers | tr -d ' ')
                local mem_usage=$(ps -p $pid -o %mem --no-headers | tr -d ' ')
                
                if (( $(echo "$cpu_usage > 50" | bc -l) )); then
                    log_warning "$process 进程CPU使用率过高: ${cpu_usage}%"
                elif (( $(echo "$mem_usage > 20" | bc -l) )); then
                    log_warning "$process 进程内存使用率过高: ${mem_usage}%"
                else
                    log_success "$process 进程运行正常 (CPU: ${cpu_usage}%, MEM: ${mem_usage}%)"
                fi
            else
                log_error "$process 进程未运行 (PID文件存在但进程不存在)"
            fi
        else
            log_warning "$process 进程PID文件不存在"
        fi
    done
}

# 检查网络连接
check_network() {
    # 检查关键端口
    local ports=("8000" "8501" "5432" "6379")
    local port_names=("API" "Web" "PostgreSQL" "Redis")
    
    for i in "${!ports[@]}"; do
        local port="${ports[$i]}"
        local name="${port_names[$i]}"
        
        if lsof -Pi :$port -sTCP:LISTEN -t >/dev/null 2>&1; then
            log_success "$name 端口 $port 正常监听"
        else
            if [ "$name" = "API" ] || [ "$name" = "Web" ]; then
                log_error "$name 端口 $port 未监听"
            else
                log_warning "$name 端口 $port 未监听"
            fi
        fi
    done
    
    # 检查外部网络连接
    if ping -c 1 -W 3 8.8.8.8 > /dev/null 2>&1; then
        log_success "外部网络连接正常"
    else
        log_warning "外部网络连接异常"
    fi
}

# 检查日志错误
check_logs() {
    if [ -d "logs" ]; then
        local error_count=0
        
        for log_file in logs/*.log; do
            if [ -f "$log_file" ]; then
                # 检查最近1小时的错误
                local recent_errors=$(find "$log_file" -mmin -60 -exec grep -i "error\|exception\|failed\|critical" {} \; 2>/dev/null | wc -l)
                
                if [ "$recent_errors" -gt 10 ]; then
                    log_error "$(basename $log_file) 最近1小时有 $recent_errors 个错误"
                    error_count=$((error_count + recent_errors))
                elif [ "$recent_errors" -gt 0 ]; then
                    log_warning "$(basename $log_file) 最近1小时有 $recent_errors 个错误"
                    error_count=$((error_count + recent_errors))
                fi
            fi
        done
        
        if [ "$error_count" -eq 0 ]; then
            log_success "日志文件无异常错误"
        fi
    else
        log_warning "日志目录不存在"
    fi
}

# 检查数据完整性
check_data_integrity() {
    python3 -c "
import sys
sys.path.append('.')
try:
    from aiquant.database.connection import DatabaseManager
    db = DatabaseManager()
    
    with db.get_connection() as conn:
        cursor = conn.cursor()
        
        # 检查关键表的数据量
        tables = ['users', 'strategies', 'orders', 'positions', 'trades']
        
        for table in tables:
            try:
                cursor.execute(f'SELECT COUNT(*) FROM {table}')
                count = cursor.fetchone()[0]
                print(f'{table}: {count} 条记录')
            except Exception as e:
                print(f'{table}: 查询失败 - {e}')
                sys.exit(1)
                
        # 检查数据一致性
        cursor.execute('''
            SELECT COUNT(*) FROM orders o 
            LEFT JOIN strategies s ON o.strategy_id = s.id 
            WHERE s.id IS NULL
        ''')
        orphan_orders = cursor.fetchone()[0]
        
        if orphan_orders > 0:
            print(f'发现 {orphan_orders} 个孤立订单记录')
            sys.exit(1)
        else:
            print('数据一致性检查通过')
            
except Exception as e:
    print(f'数据完整性检查失败: {e}')
    sys.exit(2)
" 2>/dev/null

    case $? in
        0)
            log_success "数据完整性检查通过"
            ;;
        1)
            log_warning "数据完整性存在问题"
            ;;
        2)
            log_error "数据完整性检查失败"
            ;;
    esac
}

# 生成健康报告
generate_health_report() {
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    local report_file="logs/health_report_$(date '+%Y%m%d_%H%M%S').json"
    
    if [ ! -d "logs" ]; then
        mkdir -p logs
    fi
    
    # 收集系统信息
    local cpu_usage=$(top -bn1 | grep "Cpu(s)" | awk '{print $2}' | cut -d'%' -f1 | cut -d',' -f1)
    local mem_info=$(free | grep Mem)
    local total_mem=$(echo $mem_info | awk '{print $2}')
    local used_mem=$(echo $mem_info | awk '{print $3}')
    local mem_usage=$(echo "scale=1; $used_mem * 100 / $total_mem" | bc)
    local disk_usage=$(df / | tail -1 | awk '{print $5}' | cut -d'%' -f1)
    local load_avg=$(uptime | awk -F'load average:' '{print $2}' | tr -d ' ')
    
    # 生成JSON报告
    cat > "$report_file" << EOF
{
    "timestamp": "$timestamp",
    "overall_status": "$HEALTH_STATUS",
    "system_resources": {
        "cpu_usage": $cpu_usage,
        "memory_usage": $mem_usage,
        "disk_usage": $disk_usage,
        "load_average": "$load_avg"
    },
    "services": {
        "api": $(curl -f -s http://localhost:8000/health > /dev/null 2>&1 && echo "true" || echo "false"),
        "web": $(curl -f -s http://localhost:8501 > /dev/null 2>&1 && echo "true" || echo "false"),
        "database": $(python3 -c "import sys; sys.path.append('.'); from aiquant.database.connection import DatabaseManager; DatabaseManager().test_connection()" > /dev/null 2>&1 && echo "true" || echo "false"),
        "redis": $(redis-cli ping > /dev/null 2>&1 && echo "true" || echo "false")
    },
    "network": {
        "api_port": $(lsof -Pi :8000 -sTCP:LISTEN -t >/dev/null 2>&1 && echo "true" || echo "false"),
        "web_port": $(lsof -Pi :8501 -sTCP:LISTEN -t >/dev/null 2>&1 && echo "true" || echo "false"),
        "external_connectivity": $(ping -c 1 -W 3 8.8.8.8 > /dev/null 2>&1 && echo "true" || echo "false")
    }
}
EOF
    
    log_success "健康报告已生成: $report_file"
}

# 主函数
main() {
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    
    case "${1:-}" in
        --help|-h)
            echo "AIQuant系统健康检查脚本"
            echo "用法: $0 [选项]"
            echo "选项:"
            echo "  --help, -h        显示帮助信息"
            echo "  --report          生成详细健康报告"
            echo "  --quiet           静默模式，仅输出错误"
            echo "  --json            JSON格式输出"
            exit 0
            ;;
        --quiet)
            exec > /dev/null 2>&1
            ;;
        --json)
            generate_health_report
            cat "logs/health_report_$(date '+%Y%m%d_%H%M%S').json"
            exit 0
            ;;
    esac
    
    echo "AIQuant系统健康检查 - $timestamp"
    echo "========================================"
    
    # 执行所有健康检查
    check_api_health
    check_web_health
    check_database_health
    check_redis_health
    check_system_resources
    check_processes
    check_network
    check_logs
    check_data_integrity
    
    echo ""
    echo "========================================"
    
    # 根据健康状态输出结果
    case $HEALTH_STATUS in
        0)
            log_success "系统整体健康状态良好"
            ;;
        1)
            log_warning "系统存在一些警告，建议关注"
            ;;
        2)
            log_error "系统存在严重问题，需要立即处理"
            ;;
    esac
    
    if [ "$1" = "--report" ]; then
        generate_health_report
    fi
    
    exit $HEALTH_STATUS
}

# 执行主函数
main "$@"